# !/usr/bin/env python
# -*- coding=utf-8 -*-

import numpy as np
import pandas as pd

from dataclasses import dataclass
from datetime import datetime
import time
import logging
logging.basicConfig(level=logging.NOTSET)


TICK_DATA_PATH = r"C:\Users\新田草\Desktop\正仁量化\data\splited_order.csv"
BAR_DATA_PATH = r"C:\Users\新田草\Desktop\正仁量化\data\bar_data_px_mv.csv"
BAR_DATA_PATH_TR = r"C:\Users\新田草\Desktop\正仁量化\data\bar_data_tr.csv"


def get_data(file_path, func=open) -> pd.DataFrame:
    """
    从文件中 IO 获取数据
    :param file_path:
    :param func:
    :return:
    """
    with func(file_path) as f:
        temp_data = pd.read_csv(f)
    return temp_data


# 读取数据
data = get_data(TICK_DATA_PATH)
data = data.iloc[:, 1:]
print(data.head(10))
# 将 Tick 数据按时间先后排序好
data = data.sort_values(by="Timestamp_time")


@dataclass
class BarData(object):
    """
    Candlestick bar data of a certain price movement.
    """
    contract: str
    datetime: datetime

    Timestamp: float = 0
    volume: float = 0
    open_price: float = 0
    high_price: float = 0
    low_price: float = 0
    close_price: float = 0

    bid_side_volume: int = 0
    ask_side_volume: int = 0
    bid_ask_ratio: int = 0.0


class ColsIndex(object):
    """
    列索引容器
    """

    def __init__(self):
        self.contract_ind = None
        self.Timestamp_time_ind = None
        self.Tradeprice_ind = None
        self.Tradevolume_ind = None
        self.Timestamp_ind = None
        self.direction_ind = None


def get_col_index(columns: list):
    cols_index = ColsIndex()
    for ind, each in enumerate(columns):
        if each == "Stock":
            cols_index.contract_ind = ind
        if each == "Timestamp_time":
            cols_index.Timestamp_time_ind = ind
        if each == "Timestamp":
            cols_index.Timestamp_ind = ind
        if each == "Tradeprice":
            cols_index.Tradeprice_ind = ind
        if each == "Tradevolume":
            cols_index.Tradevolume_ind = ind
        if each == "Direction":
            cols_index.direction_ind = ind

    return cols_index


data_cols = data.columns
data_cols_index = get_col_index(data_cols)


def bar_generator_tr(input_data: pd.DataFrame,
                     price_mv_ratio: float = 0.001) -> list:

    tick_datas = input_data.values

    # 委买、委卖、委比
    bid_side_volume = 0
    ask_side_volume = 0
    bid_ask_ratio = 0.0

    open_price = 0
    high_price = 0
    low_price = np.inf
    close_price = 0
    volume_trade = 0

    bar_container = list()
    pre_bar = None
    new_bar_flag = True
    for TickData in tick_datas:

        last_price = TickData[data_cols_index.Tradeprice_ind]

        direction = TickData[data_cols_index.direction_ind]
        if direction == "B":
            bid_side_volume += TickData[data_cols_index.Tradevolume_ind]
        else:
            ask_side_volume += TickData[data_cols_index.Tradevolume_ind]

        # 生成第一根 bar 逻辑
        if not pre_bar:

            if new_bar_flag:
                new_bar_flag = False
                open_price = last_price

            if high_price < last_price:
                high_price = last_price

            if low_price > last_price:
                low_price = last_price

            volume_trade += TickData[data_cols_index.Tradevolume_ind]

            if abs(high_price - low_price) >= open_price * price_mv_ratio:

                close_price = last_price

                bid_ask_ratio = (bid_side_volume - ask_side_volume) / (bid_side_volume + ask_side_volume)
                bid_ask_ratio = round(bid_ask_ratio, 3)
                # 生成一个 Bar 实例
                bar = BarData(close_price=close_price,
                              open_price=open_price,
                              low_price=low_price,
                              high_price=high_price,
                              volume=volume_trade,
                              contract=TickData[data_cols_index.contract_ind],
                              datetime=TickData[data_cols_index.Timestamp_time_ind],
                              Timestamp=TickData[data_cols_index.Timestamp_ind],
                              bid_side_volume=bid_side_volume,
                              ask_side_volume=ask_side_volume,
                              bid_ask_ratio=bid_ask_ratio)

                bar_container.append(bar.__dict__)

                open_price = last_price
                high_price = - np.inf
                low_price = np.inf
                close_price = 0.0
                volume_trade = 0

                bid_side_volume = 0
                ask_side_volume = 0
                bid_ask_ratio = 0.0

                pre_bar = bar
                new_bar_flag = True

        else:
            if new_bar_flag:
                new_bar_flag = False
                open_price = last_price

            if high_price < last_price:
                high_price = last_price

            if low_price > last_price:
                low_price = last_price

            volume_trade += TickData[data_cols_index.Tradevolume_ind]

            # 计算 true range (tr)
            price_range1 = abs(high_price - low_price)
            price_range2 = abs(pre_bar.close_price - high_price)
            price_range3 = abs(pre_bar.close_price - low_price)

            true_range = max([price_range1, price_range2, price_range3])

            if true_range >= open_price * price_mv_ratio:

                close_price = last_price

                bid_ask_ratio = (bid_side_volume - ask_side_volume) / (bid_side_volume + ask_side_volume)
                bid_ask_ratio = round(bid_ask_ratio, 3)

                # 生成一个 Bar
                bar = BarData(close_price=close_price,
                              open_price=open_price,
                              low_price=low_price,
                              high_price=high_price,
                              volume=volume_trade,
                              contract=TickData[data_cols_index.contract_ind],
                              datetime=TickData[data_cols_index.Timestamp_time_ind],
                              Timestamp=TickData[data_cols_index.Timestamp_ind],
                              bid_side_volume=bid_side_volume,
                              ask_side_volume=ask_side_volume,
                              bid_ask_ratio=bid_ask_ratio)

                bar_container.append(bar.__dict__)

                open_price = last_price
                high_price = - np.inf
                low_price = np.inf
                close_price = 0.0
                volume_trade = 0

                bid_side_volume = 0
                ask_side_volume = 0
                bid_ask_ratio = 0.0

                pre_bar = bar
                new_bar_flag = True

    return bar_container


def bar_generator_price_mv(input_data: pd.DataFrame,
                           price_mv_ratio: float = 0.001) -> list:

    tick_datas = input_data.values

    # 委买、委卖、委比
    bid_side_volume = 0
    ask_side_volume = 0
    bid_ask_ratio = 0.0

    direction = tick_datas[0][data_cols_index.direction_ind]
    if direction == "B":
        bid_side_volume += tick_datas[0][data_cols_index.Tradevolume_ind]
    else:
        ask_side_volume += tick_datas[0][data_cols_index.Tradevolume_ind]

    # 初始化 open_price
    open_price = tick_datas[0][data_cols_index.Tradeprice_ind]

    high_price = 0.0
    low_price = np.inf
    close_price = 0.0

    # 初始化 volume_trade
    volume_trade = int(tick_datas[0][data_cols_index.Tradevolume_ind])

    # 容器，装入生成的 bar
    bar_container = list()

    # 逐笔成交数据按时间先后顺序逐个推送
    for TickData in tick_datas[1:]:

        direction = TickData[data_cols_index.direction_ind]
        if direction == "B":
            bid_side_volume += tick_datas[0][data_cols_index.Tradevolume_ind]
        else:
            ask_side_volume += tick_datas[0][data_cols_index.Tradevolume_ind]

        last_price = TickData[data_cols_index.Tradeprice_ind]

        # 更新 volume_trade
        volume_trade += int(TickData[data_cols_index.Tradevolume_ind])

        # 更新 high_price / low_price
        if high_price < last_price:
            high_price = last_price
        if low_price > last_price:
            low_price = last_price

        # 价格位移达到 price_mv_ratio，生成一根 Bar
        if abs(last_price - open_price) >= open_price * price_mv_ratio:
            # 计算委比
            bid_ask_ratio = (bid_side_volume - ask_side_volume) / (bid_side_volume + ask_side_volume)
            bid_ask_ratio = round(bid_ask_ratio, 3)
            # 更新 close_price
            close_price = last_price

            # 生成一个 Bar
            bar = BarData(close_price=close_price,
                          open_price=open_price,
                          low_price=low_price,
                          high_price=high_price,
                          volume=volume_trade,
                          contract=TickData[data_cols_index.contract_ind],
                          datetime=TickData[data_cols_index.Timestamp_time_ind],
                          Timestamp=TickData[data_cols_index.Timestamp_ind],
                          bid_side_volume=bid_side_volume,
                          ask_side_volume=ask_side_volume,
                          bid_ask_ratio=bid_ask_ratio)
            # 存储新生成的Bar
            bar_container.append(bar.__dict__)

            open_price = last_price
            high_price = 0.0
            low_price = np.inf
            close_price = 0.0
            volume_trade = 0

            bid_side_volume = 0
            ask_side_volume = 0
            bid_ask_ratio = 0.0

    return bar_container   # 数据结构: [{}, {}, {}, ...]


def main():

    # 位移: price_mv_ratio
    ratio = 0.01
    # bars = bar_generator_price_mv(data, ratio)

    bars = bar_generator_tr(data, ratio)

    print("每根K线位移比例（涨跌幅）为：%s, 共合成 %s个bar." % (ratio, len(bars)))

    results = pd.DataFrame.from_dict(bars)

    results = results[["contract", "datetime", "Timestamp", "open_price",
                       "high_price", "low_price", "close_price", "volume",
                       "bid_side_volume", "ask_side_volume", "bid_ask_ratio"]]

    results.columns = ["contract", "time", "Timestamp", "open", "high",
                       "low", "close", "volume", "bid_side_vol",
                       "ask_side_vol", "bid_ask_ratio"]

    # results.to_csv(BAR_DATA_PATH)
    results.to_csv(BAR_DATA_PATH_TR)


if __name__ == '__main__':
    start_time = time.time()
    main()
    print("共耗时：%s 秒" % (time.time() - start_time))

